import os
import joblib
import numpy as np
import itertools

class ModelInference:
    def __init__(self, model_path, param_ranges, target_min, target_max,progress_bar):
        self.model_path = model_path
        self.param_ranges = param_ranges
        self.target_min = target_min
        self.target_max = target_max
        self.model = self.load_model()
        self.progress_bar = progress_bar  # 保存进度条控件
        # if hasattr(self.model, 'feature_importances_'):
        #     importances = self.model.feature_importances_
        #     feature_importance = dict(zip(self.param_ranges.keys(), importances))
        #     print("特征重要性:", feature_importance)

    # 加载模型函数
    def load_model(self):
        if os.path.exists(self.model_path):
            print(f"模型文件存在，文件大小为: {os.path.getsize(self.model_path)} 字节")
            with open(self.model_path, 'rb') as f:
                model = joblib.load(f)  # 确保这里返回的是模型对象
            return model
        else:
            print("模型文件不存在")
            return None

    # 定义推理函数
    def predict_with_model(self, new_data):
        """
        使用训练好的模型对新的数据进行预测
        :param new_data: 需要推理的数据，格式为2D数组（n_samples, n_features）
        :return: 预测结果
        """
        predictions = self.model.predict(new_data).ravel()  # 确保model是已训练的模型对象  ravel将二维数组转换为一维
        return predictions

    # 生成参数值列表，处理步长、最大值、最小值的逻辑
    def generate_values(self, min_val, max_val, step):
        # 步长不允许为负数
        if step is not None and step < 0:
            raise ValueError("步长不允许为负数")

        # 最小值不允许比最大值大
        if min_val > max_val:
            raise ValueError("最小值不能大于最大值")

        # 步长为空值、0或大于最大值与最小值之差时，只返回最小值和最大值
        if step is None or step == 0 or step > (max_val - min_val):
            if min_val == max_val:
                return [round(min_val, 3)]  # 最大值等于最小值时，只返回一个值
            else:
                return [round(min_val, 3), round(max_val, 3)]

        # 计算包含最大值的步长序列
        values = np.arange(min_val, max_val, step).tolist()
        # 四舍五入每个生成的值
        values = [round(val, 3) for val in values]

        if round(values[-1], 3) != round(max_val, 3):
            values.append(round(max_val, 3))  # 确保最大值被包含
        return values

    # 生成属性的所有组合
    def generate_combinations(self):
        param_names = list(self.param_ranges.keys())
        param_values = [self.generate_values(start, stop, step) for start, stop, step in self.param_ranges.values()]
        combinations = list(itertools.product(*param_values))
        return param_names, combinations

    # 筛选符合目标范围的结果
    def filter_predictions(self, combinations):
        filtered_results = []
        total_combinations = len(combinations)
        self.progress_bar.setMaximum(total_combinations)  # 设置进度条的最大值为总组合数

        for idx, combination in enumerate(combinations):
            new_data = np.array([combination])
            prediction = self.predict_with_model(new_data)

            if self.target_min <= prediction <= self.target_max:
                filtered_results.append((combination, prediction))

            # 更新进度条
            self.progress_bar.setValue(idx + 1)

        return filtered_results

    # 主推理函数
    def run_inference(self):
        if not self.model:
            print("模型未加载，推理无法进行")
            return {}

        # 生成所有属性的组合
        param_names, combinations = self.generate_combinations()
        print(f"生成的组合数量: {len(combinations)}")
        # print(combinations)
        # 对每个组合进行预测并筛选
        filtered_results = self.filter_predictions(combinations)
        # print(filtered_results)
        # 保存符合条件的组合和预测值到字典
        results_dict = {}
        for i, (combination, prediction) in enumerate(filtered_results):
            results_dict[f'comb_{i}'] = {'combination': combination, 'target': prediction}

        # 返回字典
        return results_dict

# 使用类
if __name__ == "__main__":
    # model_path = '../main/Data_Management_page_main/model/20241014140423_gb_model.pkl'

    # model_path = '../main/Data_Management_page_main/models/20241030090129/20241030090140_rid_model.pkl'
    # model_path = '../main/Data_Management_page_main/models/20241030090129/20241030090140_svm_model.pkl'
    model_path = '../main/Data_Management_page_main/models/20241101081849/20241101081900_gb_model.pkl'

    param_ranges = {
        'Si/％': (0.06, 0.06, 1),
        'Mn/％': (0.5, 0.5, 1),
        'Zn/％': (5.8, 5.8, 1),
        'Mg/％': (2.17, 2.17, 1),
        'Cu/％': (1.7, 1.7, 1),
        'Cr/％': (0.10, 100.00, 0.01),
        'Zr/％': (0, 0, 1),
        'Ti/％': (0.019, 0.019, 1),
        'Fe/％': (0.086, 0.086, 1),
        'Ni/％': (0.025, 0.025, 1),
        'other': (0, 0, 1),
    }
    target_min = 0
    target_max = 10000

    # 创建类的实例
    inference = ModelInference(model_path, param_ranges, target_min, target_max,None)

    # 运行推理
    inference.run_inference()
